library(ggpubr)
library(Rsamtools)
library(tidyverse)

ah.meta <- read_delim('mission/ye_AH-cOME/ah-come.meta.txt')

# try annotate cell with G/R from SAM ------------
call.g2r <- function(bam, coord = 'chr14:105737776',
                     mapq = 10, debug = FALSE, verbose = FALSE) {
  pos.on.chr <- str_extract(coord, '\\d+$') |> as.integer()
  
  bam.param <- ScanBamParam(tag = c('CB'), mapqFilter = mapq,
                            which = GRanges(coord),
                            what = scanBamWhat())
  
  scaned <- scanBam(bam, param = bam.param) |>
    pluck(1)
  
  scaned$tag <- scaned$tag$CB
  scaned$seq <- scaned$seq |> as.character()
  scaned$qual <- scaned$qual |> as.character()
  
  scaned.tidy <- scaned |>
    as_tibble() |>
    filter(tag != '-') |>
    mutate(left.clip = str_extract(cigar, '^\\d+(?=S)') |> as.integer() |> replace_na(0),
           right.clip = str_extract(cigar, '\\d+(?=S$)') |> as.integer() |> replace_na(0),
           snp.pos = pos.on.chr + 1 - pos + left.clip,
           snp.base = str_sub(seq, snp.pos, snp.pos),
           baseq = str_sub(qual, snp.pos, snp.pos) |> sapply(charToRaw) |> as.integer() - 33L,
           mapq, tag, qname,
           .keep = 'none')
  
  if(verbose) scaned.tidy |> dplyr::count(mapq >= 30, snp.base) |> print()
  
  allele.included <- scaned.tidy |>
    summarise(n = n(), .by = tag) |>
    filter(n > 1) |>
    left_join(scaned.tidy) |>
    summarise(base.n = n(), .by = c(tag, snp.base, n)) |>
    filter(n > base.n)
  
  if (nrow(allele.included) > 0) {
    warning('Found allele included barcodes!')
    scaned.tidy <- scaned.tidy |>
      mutate(snp.base = ifelse(tag %in% allele.included$tag, 
                               str_c(snp.base, '_X'), snp.base))
  }
  
  scaned.tidy |>
    mutate(.cell = tag, g2r.base = snp.base, baseq, mapq,
           .keep = ifelse(debug, 'unused', 'none'), .before = 1)
}

call.g2r('~/append-ssd/nextflowing/ye.ah.come1/P01cb/P01.recal.bam') |>
  filter(str_ends(g2r.base, 'X'))

call.g2r('~/append-ssd/nextflowing/ye.ah.come1/P03cb/P03.recal.bam', debug = T)

yeah.path <- '~/append-ssd/nextflowing/ye.ah.come1'

recaled.path <- list.files(yeah.path, 'recal.bam$', full.names = T, recursive = T)

bc.g2r.call <- recaled.path |>
  str_extract('P0[1-9]') |>
  set_names(x = recaled.path, nm = _) |>
  map(call.g2r, .progress = T) |>
  list_rbind(names_to = 'pid')

bc.g2r.call |> filter(str_detect(g2r.base, 'X'))

bc.g2r.call |>
  mutate(gbc = str_c(pid, '_', .cell)) |>
  slice_max(baseq,by = gbc) |>
  ggplot(aes(baseq)) + geom_dotplot()

bc.g2r.call |>
  write_csv('mission/ye_AH-cOME/ah-come.g2r.bc.csv')

# compare with gatk vcf --------
i2t.call <-
  read_delim('~/append-ssd/nextflowing/ye.ah.come1/empty_intron/rnavar_out/cllrg.i2t.txt',
             comment = '##') |>
  select(10:last_col()) |>
  pivot_longer(everything(), names_to = 'sample', values_to = 'gt') |>
  separate_wider_delim(gt, ':', names_sep = '.')

i2t.call |>
  mutate(id = sample, i2t.gt = gt.1, i2t.base.count = gt.2, .keep = 'none') |>
  right_join(ah.meta) |>
  arrange(pid) |>
  write_csv('yeah.come.i2t.csv')

shen21escc <- '/home/supervisor/mist2/gjsx/shen2020escc/test5p/' |>
  list.files('bam', full.names = T, recursive = T)

tibble(bam = shen21escc,
       sample = str_extract(bam, 'Sample_\\d+')) |>
  relocate(sample) |>
  write_csv('/home/supervisor/mist2/gjsx/shen2020escc/nxf.bam.csv')
